from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

class Reranker:
    def __init__(self,) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained('/hy-tmp/bge-reranker-large')
        self.model = AutoModelForSequenceClassification.from_pretrained('/hy-tmp/bge-reranker-large').to_bettertransformer()
        self.model.eval()
        self.model.cuda()
        print(r'成功加载了reranker-BAAI/bge-reranker-large-------------------------------------------------------------------------------------------------------')
        print('==========================================================================================================================================================')

    def rank(self,input,candidates):
        pairs=[[input,c] for c in candidates]
        with torch.inference_mode():
            inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512).to('cuda')
            scores = self.model(**inputs, return_dict=True).logits.view(-1, ).float()
        sorted_combined = sorted(list(zip(scores, range(len(scores)))), key=lambda x: x[0], reverse=True)  # 根据分数（第一个元素）排序
        sorted_items = [item for _, item in sorted_combined]
        return [candidates[i] for i in sorted_items]



# scores = [85, 70, 90, 60]
# items = ['A', 'B', 'C', 'D']

# s=[['a',1],['a',2],['a',3],['a',4],['a',5]]
# sorted_combined = sorted(list(zip(scores, range(len(scores)))), key=lambda x: x[0], reverse=True)  # 根据分数（第一个元素）排序
# sorted_items = [item for _, item in sorted_combined]


# print(sorted_items)  # 输出排序后的项目名列表

# print(s[sorted_items[:3]]) 

if __name__=='__main__':
    r=Reranker()
    candidates=['最低分数','最低位次','录取人数']
    import pandas as pd
    # 读取整个 Excel 文件
    df = pd.read_excel('/home/lxy/multiR/分类测试1.xlsx')
    
    # 选择特定行范围（例如，从第1行到第5行）
    specific_rows = df.iloc[0:5]
    # 选择特定列范围（例如，从第A列到第C列）
    specific_columns = df.loc[:,'问题']
    answer_columns=df.loc[:,'期待的分类']
    wrong_cnt=0
    for i,question in enumerate(specific_columns):
        fake_answer=r.rank(question,candidates)[0]
        if fake_answer!=answer_columns[i]:
            wrong_cnt+=1
            print(question,f'输出答案{fake_answer}',f'期待答案{answer_columns[i]}')
    print(f'{wrong_cnt}/{len(question)}')